MAX_DEGREE = 2

import copy
import sys
import random
import logging
import optparse
import networkx as nx

_logger = logging.getLogger("dance")

from eap import algorithms
from eap import base
from eap import creator
from eap import halloffame
from eap import toolbox

import graph as dgraph

def main(options, args):
    logging.basicConfig(level=getattr(logging, options.loglevel), stream=sys.stderr)
    graph = dgraph.parseData(args[0], args[1])
    tools = toolsInit(options, graph)

    pop = tools.population()
    hof = halloffame.ParetoFront()
    algorithms.eaMuPlusLambda(tools, pop, options.mu, options.lambd,
            options.cxpb, options.mutpb, options.ngen, hof)
    for good in tools.select(pop,4):
    #for good in hof: 
        print good.fitness.values
        if getattr(logging, options.loglevel) >= logging.INFO:
            printPartnerships(good)
    #unique = set([tuple(ind.edges()) for ind in hof])
    #for ind in unique:
    #    printPartnershipsFromEdgelist(ind, graph)
    #print len(hof), len(unique)

def splitGraph(graph):
    """
    Split a set of nodes into odds and evens, return them ordered by size.
    """
    nodes = graph.nodes()
    evens = set(filter((lambda x: x % 2 ^ 1), nodes))
    odds = set(filter((lambda x: x % 2), nodes))
    return sorted([odds, evens], cmp=(lambda x,y: cmp(len(x), len(y))))

def graphInd(graph, fitness):
    # Doesn't do anything in particular, just a workaround some bug in DEAP
    new = graph()
    setattr(new, 'fitness', fitness())
    return new

def toolsInit(options, graph):
    """
    Set up the toolbox based on provided options
    """
    creator.create("Fitness", base.Fitness, weights=(1.0, 1.0, 1.0, 1.0))
    creator.create("Individual", nx.Graph, fitness=creator.Fitness)
    creator.create("Population", list)
    left, right = splitGraph(graph)

    tools = toolbox.Toolbox()

    # Structure initializers
    tools.register("attr_bipartite", randomBipartite, graph, 2)
    tools.register("individual", graphInd,
        tools.attr_bipartite, creator.Fitness) 
    tools.register("population", creator.Population,
        content_init=tools.individual, size_init=options.popSize)
    tools.register("info", printPartnerships)
    
    # Operator registering

    tools.register("select", getattr(toolbox, options.select))
    tools.register("mate", toolbox.cxTwoPoints) # Irrelevant for now, I think
    tools.register("mutate", mutDancePartners, graph, 1.0/len(left))

    tools.register("evaluate", evalPartnering)

    return tools

def randomBipartite(graph, maxDegree):
    """
    Generates a random bipartite graph based on a set of nodes,
    with even nodes on one side and odd nodes on the other.
    All nodes have degree at least 1, until they don't.
    """
    small, big = splitGraph(graph)
    bipartite = nx.Graph(graph)
    bipartite.remove_edges_from(bipartite.edges())
    for right in big:
        for left in small:
            if random.random() < ((float(len(small))/len(big))**-1)/len(big):
                bipartite.add_edge(left, right, graph[left][right])

    enforceDegree(graph, bipartite)

    return bipartite

def worstPair(ind, p):
    #return sorted(ind[p].items(), cmp=(lambda x,y: cmp(ind[p][x[0]]['weight'],
    #    ind[p][y[0]]['weight'])))[0][0]
    return random.choice(ind[p].keys())

def enforceDegree(graph, ind, maxDegree=MAX_DEGREE):
    """
    Tries to force all nodes to have degree 0 < d <= 2.
    Operates first on the larger side of the bipartite graph,
    then on the smaller side. Operations on the smaller side
    often reduce the degree of some node on the larger side to zero,
    but there's little that can be done about that.
    """
    left, right = splitGraph(graph)
    for A,B in ((right, left), (left, right)):
        for a in A:
            while True:
                if ind.degree(a) < 1:
                    pot = ind.node[a]['untried']
                    if not pot:
                        ind.node[a]['untried'] = sorted(graph[a].keys(),
                                cmp=(lambda x,y: cmp(graph[a][x], graph[a][y])))
                        pot = ind.node[a]['untried']
                        random.shuffle(pot)
                    b = pot.pop()
                    ind.add_edge(a, b, graph[a][b])
                    break
                if ind.degree(a) > maxDegree:
                    b = worstPair(ind, a)
                    ind.remove_edge(a, b)
                    continue
                break

def mutDancePartners(graph, indpb, ind):
    """
    Add/remove pairings at random. Strictly enforces degree range.
    """
    left, right = splitGraph(ind)
    old = ind.edges()
    for p in left:
        if random.random() < indpb:
            cur = ind[p].keys() # current partnerships
            pot = ind.node[p]['untried']
            if not pot:
                ind.node[p]['untried'] = sorted(graph[p].keys(),
                        cmp=(lambda x,y: cmp(graph[p][x], graph[p][y])))
                pot = ind.node[p]['untried']
                random.shuffle(pot)
                #_logger.info(pot)
            if random.random() < 1:
                new = pot.pop()
                ind.add_edge(p, new, graph[p][new])
            if random.random() < 1 and cur:
                worst = worstPair(ind, p)
                ind.remove_edge(p, worst)
    enforceDegree(graph, ind)
    new =  evalPartnering(ind)
    return ind,

def evalPartnering(ind):
    """
    Evaluate pairings described by ind on the following:
    - How well preferences are satisfied
    - How well heights are matched
    - How well weighted preferences (preferences plus
    information about attendance, seniority, etc.)
    - How few people are not partnered at all
    - Prefers to partner the smaller side with as few
    partners as possible, weakly, as part of the two
    preference-based metrics.
    """
    return evalPreferences(ind), evalHeight(ind), \
            evalPreferencesWeighted(ind), evalOrphans(ind)
    #return 1.0, evalHeight(ind), 1.0, evalOrphans(ind) 

def evalOrphans(ind):
    """
    Fitness worsens as more and more people are unpaired.
    """
    orphans = 0.0
    for p in ind.nodes():
        if not ind.degree(p):
            orphans += 1
    return -(2**orphans)

def evalPreferences(ind):
    """
    Fitness improves as people are paired with
    others they want to be paired with.
    """
    totalSatisfaction = 0.0
    for p,q in ind.edges_iter():
        totalSatisfaction += ind[p][q]['weight']
    try:
        return 2**(totalSatisfaction/totalPartnerships(ind))
    except ZeroDivisionError: return 0

def evalHeight(ind):
    """
    Fitness improves as people are paired with
    others who are close to them in height.
    """
    discrepancy = 0
    for p,q in ind.edges_iter():
        discrepancy += abs(ind.node[p]['height'] - ind.node[q]['height'])
    try:
        return -(discrepancy/totalPartnerships(ind))
    except ZeroDivisionError: return -1.0

def evalPreferencesWeighted(ind):
    """
    Fitness improves similarly to evalPreferences,
    but with preferences weighted based on a dancer's
    importance, based on several factors.
    """
    totalSatisfaction = 0.0
    for p,q in ind.edges_iter():
        pmult = ind.node[p]['attendance'] * (ind.node[p]['seniority']**1.3 +
                ind.node[p]['sadness'])
        qmult = ind.node[q]['attendance'] * (ind.node[q]['seniority']**1.3 +
                ind.node[q]['sadness'])
        if p % 2: lmult, fmult = qmult, pmult
        else: lmult, fmult = qmult, pmult
        totalSatisfaction += ind[p][q]['lead'] * lmult
        totalSatisfaction += ind[p][q]['follow'] * fmult
    try:
        return 2**(totalSatisfaction/totalPartnerships(ind))
    except ZeroDivisionError: return 0

def totalPartnerships(ind):
    return len(ind.edges())

def printPartnerships(ind):
    """
    Prints all partnerships in a given graph
    """
    left, right = splitGraph(ind)
    for p in left:
        sys.stderr.write(ind.node[p]['name']+':')
        for q in ind[p]:
            sys.stderr.write('\t' + ind.node[q]['name'])
        print
    sys.stderr.write("UNPAIRED:")
    for q in right:
        if not ind.degree(q):
            sys.stderr.write('\t' + ind.node[q]['name'])
    print

def printPartnershipsFromEdgelist(ind, graph):
    """
    Like above, but regenerates the individual
    from the top-level graph and an edge list.
    """
    edges = ind
    ind = nx.Graph(graph)
    ind.remove_edges_from(graph.edges())
    ind.add_edges_from(edges)
    left, right = splitGraph(ind)
    for p in left:
        sys.stderr.write(ind.node[p]['name']+':')
        for q in ind[p]:
            sys.stdout.write('\t' + ind.node[q]['name'])
        print
    sys.stdout.write("UNPAIRED:")
    for q in right:
        if not ind.degree(q):
            sys.stdout.write('\t' + ind.node[q]['name'])
    print


def getOptions():
    """
    Parses arguments specified from the commandline
    """
    parser = optparse.OptionParser()

    parser.add_option("--nsga2", action="store_const", const='nsga2', dest='select',
            help='Use NSGA2 instead of SPEA2')
    parser.add_option("-m", "--mu", action="store", type="int", dest='mu', default=30)
    parser.add_option("-l", "--lambda", action="store", type="int", dest='lambd',
            default=60)
    parser.add_option('-c', '--cx', action='store', type='float', dest='cxpb',
            default=0, help='Crossover probability')
    parser.add_option('-t', '--mut', action='store', type='float', dest='mutpb',
            default=1, help='Mutation probability')
    parser.add_option('-n', '--ngen', action='store', type='int', dest='ngen',
            default=300, help='Number of generations')
    parser.add_option('-s', '--population', action='store', type='int', dest='popSize',
            default=60, help='Population size')
    parser.add_option('-v', '--verbosity', action='store', type='string', dest='loglevel',
            default='INFO', help='Verbosity level')

    options, args = parser.parse_args()
    if len(args) != 2:
        print "Wrong number of input files. Need 2, preference data and dancer data"
        sys.exit(1)
    if not options.select: options.select = 'spea2' 
    return options, args

if __name__ == "__main__":
    random.seed()         # Seed does not include item creation
    main(*getOptions())
